#!/usr/bin/env python3
import os
import sys
import argparse
import subprocess
import re
from pathlib import Path
import pwd
import fnmatch

def check_sudo():
    """检查是否以sudo命令或者root用户运行"""
    if os.geteuid() != 0:
        print("错误：请使用sudo命令或者root用户运行此脚本！", file=sys.stderr)
        sys.exit(1)

def get_uid_gid(path):
    """获取路径的UID和GID"""
    try:
        result = subprocess.run(
            ["ls", "-nd", path],
            capture_output=True,
            text=True,
            check=True
        )
        output = result.stdout.strip()
        parts = output.split()
        if len(parts) < 4:
            raise ValueError(f"无法解析ls输出: {output}")
        return int(parts[2]), int(parts[3])
    except subprocess.CalledProcessError as e:
        print(f"执行ls命令出错: {e.stderr}", file=sys.stderr)
        sys.exit(1)
    except ValueError as e:
        print(str(e), file=sys.stderr)
        sys.exit(1)

def get_user_uid_gid(username):
    """获取用户的UID和GID"""
    try:
        uid = subprocess.run(
            ["id", "-u", username],
            capture_output=True,
            text=True,
            check=True
        ).stdout.strip()
        
        gid = subprocess.run(
            ["id", "-g", username],
            capture_output=True,
            text=True,
            check=True
        ).stdout.strip()
        
        return int(uid), int(gid)
    except subprocess.CalledProcessError as e:
        print(f"获取用户ID失败: {e.stderr}", file=sys.stderr)
        sys.exit(1)

def find_backup_dir():
    """寻找可用的备份目录"""
    base_name = "backup_target_home"
    index = 1
    while True:
        backup_dir = f"{base_name}_{index}"
        if not os.path.exists(backup_dir):
            return backup_dir
        index += 1

def backup_home(target_home, backup_dir):
    """备份目标用户的家目录"""
    print(f"备份目标家目录到: {backup_dir}")
    try:
        subprocess.run(
            ["rsync", "-a", target_home + "/", backup_dir + "/"],
            check=True
        )
        print("备份完成！")
    except subprocess.CalledProcessError as e:
        print(f"备份失败: {e}", file=sys.stderr)
        sys.exit(1)

def calculate_size(source_home, exclude_patterns):
    """计算排除指定模式后的源目录大小（KB）"""
    try:
        # 构建排除参数
        exclude_args = []
        for pattern in exclude_patterns:
            exclude_args.extend(["--exclude", pattern])
        
        # 使用rsync模拟操作
        cmd = [
            "rsync", "-a", "--dry-run", "--stats"
        ] + exclude_args + [
            f"{source_home}/",
            "/tmp/dest/"
        ]
        
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True
        )
        
        # 改进的正则表达式：允许数字中的逗号分隔符
        match = re.search(r"Total file size:\s*([\d,]+)\s+bytes", result.stdout)
        if not match:
            print("错误：无法解析文件大小，rsync输出:\n" + result.stdout, file=sys.stderr)
            sys.exit(1)
        
        # 移除逗号后转换为整数
        total_bytes = int(match.group(1).replace(',', ''))
        return (total_bytes + 1023) // 1024  # 字节→KB（向上取整）
    
    except subprocess.CalledProcessError as e:
        print(f"Rsync执行失败 [代码 {e.returncode}]:\n{e.stderr}", file=sys.stderr)
        sys.exit(1)

def get_free_space(path):
    """获取路径所在分区的可用空间（KB）"""
    try:
        result = subprocess.run(
            ["df", "-kP", path],
            capture_output=True,
            text=True,
            check=True
        )
        lines = result.stdout.splitlines()
        if len(lines) < 2:
            raise ValueError("df命令输出格式异常")
        
        parts = lines[1].split()
        if len(parts) < 4:
            raise ValueError("df命令输出格式异常")
        
        return int(parts[3])
    except (subprocess.CalledProcessError, ValueError) as e:
        print(f"获取磁盘空间失败: {e}", file=sys.stderr)
        sys.exit(1)

def list_top_level(source_home, exclude=[]):
    top_level_dirs_dict = {}
    top_level_files_dict = {}

    # 遍历文件
    for root, dirs, files in os.walk(source_home):

        for v in dirs[:]:
            
            vpath = os.path.join(root, v) 
            
            if not vpath.startswith(source_home + os.sep):
                print(f"错误：路径{vpath}居然不以{source_home}开头")
                sys.exit(1)
            
            vpcut = vpath[len(source_home):]
        
            # 通过glob排除文件
            for vexclude in exclude:
                if fnmatch.fnmatch(vpcut, vexclude):
                    dirs.remove(v)
                    break
        
        # 判断 root 是不是 source_home
        # 如果是，则当前目录为 source_home 下的一层目录，也就是顶级目录，也就是/
        # 如果是，就需要将 dirs 拼接 root 并添加到 top_level_dirs_dict
        # 如果不是，则忽略
        # 因为之前的排除项逻辑中，对被排除的目录执行了dirs.remove()，所以被排除的目录不会记录到顶级目录
        if root == source_home:
            for v in dirs:
                top_level_dirs_dict[os.path.join(root, v)] = 0

        for v in files:
            
            vpath = os.path.join(root, v) 
            
            if not vpath.startswith(source_home + os.sep):
                print(f"错误：路径{vpath}居然不以{source_home}开头")
                sys.exit(1)
            
            vpcut = vpath[len(source_home):]
            
            # 通过glob排除文件
            is_exclude = False
            for vexclude in exclude:
                if fnmatch.fnmatch(vpcut, vexclude):
                    is_exclude = True
                    break
            
            if is_exclude:
                continue
            
            # 判断是不是顶层目录的文件，逻辑为：
            # 判断 root 是不是 source_home
            # 如果是，那就表明是顶层目录的文件
            # 如果不是，则不进行这个处理
            # 如果是，则需要获取它的大小并将其添加到 top_level_files_dict
            # 然后，避免对它进行“寻找所属顶层目录”的逻辑
            if root == source_home:
                try:
                    size_bytes = 0
                    
                    if os.path.islink(vpath):
                        size_bytes = os.lstat(vpath).st_size
                    else:
                        size_bytes = os.path.getsize(vpath)
                    
                    top_level_files_dict[vpath] = size_bytes
                    continue
                except OSError as e:
                    print(f"无法访问 {vpath}: {e}")
                    continue

            # 判断属于哪一个顶层目录
            # 将 被判断的路径 和 顶层文件夹路径 都裁剪成以source_home为/的形式
            # ps: vpcut 存在，可复用
            # 判断 被判断的路径 以哪个 顶层文件夹路径 开始，哪个就是它的顶层文件夹
            # 增加大小
            for k, v_count in top_level_dirs_dict.items():
                
                if not k.startswith(source_home + os.sep):
                    print(f"错误：路径{k}居然不以{source_home}开头")
                    sys.exit(1)
            
                kcut = k[len(source_home):]
                
                if vpcut.startswith(kcut + os.sep):
                    try:
                        size_bytes = 0
                    
                        if os.path.islink(vpath):
                            size_bytes = os.lstat(vpath).st_size
                        else:
                            size_bytes = os.path.getsize(vpath)
                    
                        top_level_dirs_dict[k]+=size_bytes
                        break
                    except OSError as e:
                        print(f"无法访问 {vpath}: {e}")
                        break
   
    # 合并
    top_level_all_dict = top_level_dirs_dict | top_level_files_dict
    # 根据大小降序
    top_level_all_list = sorted(top_level_all_dict.items(), key=lambda x: x[1], reverse=True)

    # 转换大小、生成单位并打印
    print("源家目录下的 文件/文件夹 列表：")
    for v in top_level_all_list:
        # 单位
        unit = ""
        # 数值
        num = 0
        # 累计的bytes
        size = v[1]

        if size < 1024:
            unit, num = "B", size
        elif size < 1024**2:
            unit, num = "KB", size/1024
        elif size < 1024**3:
            unit, num = "MB", size/(1024**2)
        else:
            unit, num = "GB", size/(1024**3)

        print(f"{num:.2f} {unit} {v[0]}")


def migrate_data(source_home, target_home, exclude_patterns, is_copy):
    """执行数据迁移"""
    try:
        # 构建rsync命令
        exclude_args = []
        for pattern in exclude_patterns:
            exclude_args.extend(["--exclude", pattern])
        
        cmd = ["rsync", "-a"] + exclude_args
        
        if not is_copy:
            cmd.append("--remove-source-files")
        
        cmd.extend([f"{source_home}/", f"{target_home}/"])
        
        print("开始迁移数据...")
        print("在操作完成之前不要按下Ctrl+C也不要关闭终端！")
        subprocess.run(cmd, check=True)
        print("迁移完成！")
    except subprocess.CalledProcessError as e:
        print(f"迁移失败: {e}", file=sys.stderr)
        sys.exit(1)

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="用户家目录迁移工具")
    parser.add_argument("-e", "--exclude", action="append", default=[], 
                        help="排除目录模式（rsync --exclude的语法）（假如你要排除源`家目录/aaa`，你应该写为`/aaa`），可多次指定")
    parser.add_argument("-s", "--source", required=True, 
                        help="源家目录路径（必须）（因为我要拼接/，所以你不要以/结尾）")
    parser.add_argument("-u", "--user", required=True, 
                        help="目标用户名（必须）")
    parser.add_argument("-c", "--copy", action="store_true", 
                        help="以复制形式迁移（默认为移动）")
    parser.add_argument("-nb", "--no-backup-src-home", action="store_true",
                        help="不备份源家目录（默认会备份）")
    
    args = parser.parse_args()
  
    #检查执行该命令的终端所在的工作目录
    if os.getcwd() != "/root/migrate_workdir":
        print(f"当前工作目录为{os.getcwd()}，这个目录不行，必须是/root/migrate_workdir")
        sys.exit(1)

    # 获取当前脚本的绝对路径
    if str(Path(__file__).resolve().parent) != "/root/migrate_workdir":        
        print(f"当前脚本文件所在的目录路径为{Path(__file__).resolve().parent}，这个目录不行，必须是/root/migrate_workdir")
        sys.exit(1)

    # 检查路径，要求末尾不带/
    if args.source.endswith("/"):
        print(f"错误：路径 {args.source} 因为我要拼接/，所以你不要以/结尾")
        sys.exit(1)
    
    # 步骤1: 检查是否以sudo命令或者root用户运行
    check_sudo()
    
    # 步骤2: 检查UID/GID是否匹配
    src_uid, src_gid = get_uid_gid(args.source)
    tgt_uid, tgt_gid = get_user_uid_gid(args.user)
    
    if (src_uid, src_gid) != (tgt_uid, tgt_gid):
        print(f"错误: UID/GID 不匹配！", file=sys.stderr)
        print(f"源家目录: UID={src_uid}, GID={src_gid}", file=sys.stderr)
        print(f"目标用户: UID={tgt_uid}, GID={tgt_gid}", file=sys.stderr)
        sys.exit(1)
    
    # 获取目标用户家目录
    try:
        target_home = pwd.getpwnam(args.user).pw_dir
    except KeyError:
        print(f"错误: 用户 '{args.user}' 不存在！", file=sys.stderr)
        sys.exit(1)
    
    # 步骤3: 创建备份
    backup_dir = find_backup_dir()
    if not args.no_backup_src_home:
        print(f"即将备份**目标用户**的家目录{target_home}到{backup_dir}")
        input("按回车键继续，或按Ctrl+C取消...")
        os.makedirs(backup_dir, exist_ok=True)
        backup_home(target_home, backup_dir)
    
    # 准备排除列表（添加额外排除项）
    # /.config/deepin 存储了 deepin 自带的应用程序的设置
    # /.config/systemd/user 含有一些会导致音频服务挂掉的软链接
    exclude_patterns = args.exclude + ["/.config/deepin", "/.config/systemd/user"]
    
    # 步骤4: 检查存储空间
    source_size_kb = calculate_size(args.source, exclude_patterns)
    free_space_kb = get_free_space(target_home)
    
    if source_size_kb > free_space_kb:
        print(f"错误: 磁盘空间不足！", file=sys.stderr)
        print(f"需要空间: {source_size_kb / 1024 / 1024} GB", file=sys.stderr)
        print(f"可用空间: {free_space_kb / 1024 / 1024} GB", file=sys.stderr)
        list_top_level(args.source, exclude_patterns)
        sys.exit(1)
    
    # 步骤5: 用户确认
    print("\n即将执行迁移操作，请确认以下信息:")
    print(f"额外排除目录: /.config/deepin, /.config/systemd/user")
    print(f"排除目录: {args.exclude or '无'}")
    print(f"源家目录: {args.source}")
    print(f"目标用户: {args.user}")
    if not args.no_backup_src_home:
        print(f"目标用户家目录已备份到：{backup_dir}")
    print(f"迁移模式: {'复制' if args.copy else '移动'}")
    input("按回车键继续，或按Ctrl+C取消...")
    
    # 步骤6: 执行迁移
    migrate_data(args.source, target_home, exclude_patterns, args.copy)

if __name__ == "__main__":
    main()